from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import sys
import os

class CUDAExtension(Extension):
    def __init__(self, name, sources, **kwargs):
        super().__init__(name, sources, **kwargs)

class CUDACompiler(build_ext):
    def build_extension(self, ext):        
        super().build_extension(ext)

if __name__ == "__main__":
    ext_modules = [
        CUDAExtension(
            name='my_cuda_ops',
            sources=[
                'src/bindings.cpp',
            ],
            include_dirs=[],
            language='c++',
            extra_compile_args=[
                '-O3',
                '-std=c++11',
            ],
            extra_link_args=[
                '-lmy_cuda_ops',
            ],
        )
    ]

    setup(
        name='my_cuda_ops',
        ext_modules=ext_modules,
        cmdclass={'build_ext': CUDACompiler},
        zip_safe=False,
        install_requires=['pybind11>=2.6',]
    )
